import logging
import concurrent.futures
import threading
from typing import Optional, Dict, Any

from log_config import setup_logging

# 配置日志
setup_logging()
logger = logging.getLogger(__name__)

class ThreadPoolManager:
    """
    线程池管理器 - 封装线程池的创建、配置和关闭功能
    提供线程池的优雅管理和资源控制
    """
    def __init__(self, max_workers: int = 5, thread_name_prefix: str = "rag-"):
        """
        初始化线程池管理器
        
        Args:
            max_workers: 线程池中的最大工作线程数
            thread_name_prefix: 线程名称前缀，用于日志和调试
        """
        self.max_workers = max_workers
        self.thread_name_prefix = thread_name_prefix
        self.thread_pool: Optional[concurrent.futures.ThreadPoolExecutor] = None
        self._lock = threading.RLock()
        self._closed = False
        
    def initialize(self):
        """
        初始化线程池
        
        Returns:
            ThreadPoolManager: 返回自身，支持链式调用
        
        Raises:
            RuntimeError: 如果线程池已初始化或已关闭
        """
        with self._lock:
            if self.thread_pool is not None:
                logger.warning("线程池已经初始化，无需重复初始化")
                return self
            
            if self._closed:
                raise RuntimeError("线程池已关闭，无法再次初始化")
            
            logger.info(f"初始化线程池，最大工作线程数: {self.max_workers}")
            self.thread_pool = concurrent.futures.ThreadPoolExecutor(
                max_workers=self.max_workers,
                thread_name_prefix=self.thread_name_prefix
            )
            
            return self
            
    def get_thread_pool(self) -> concurrent.futures.ThreadPoolExecutor:
        """
        获取线程池实例
        
        Returns:
            ThreadPoolExecutor: 线程池实例
        
        Raises:
            RuntimeError: 如果线程池未初始化或已关闭
        """
        with self._lock:
            if self.thread_pool is None:
                raise RuntimeError("线程池未初始化")
            
            if self._closed:
                raise RuntimeError("线程池已关闭")
            
            return self.thread_pool
            
    def shutdown(self, wait: bool = True, cancel_futures: bool = False):
        """
        关闭线程池，释放资源
        
        Args:
            wait: 是否等待所有提交的任务完成
            cancel_futures: 是否取消尚未开始运行的任务
        """
        with self._lock:
            if self.thread_pool is not None and not self._closed:
                logger.info(f"关闭线程池，等待任务完成: {wait}")
                try:
                    self.thread_pool.shutdown(
                        wait=wait,
                        cancel_futures=cancel_futures
                    )
                    logger.info("线程池已成功关闭")
                except Exception as e:
                    logger.error(f"关闭线程池时出错: {str(e)}")
                finally:
                    self.thread_pool = None
                    self._closed = True
            else:
                logger.warning("线程池已关闭或未初始化，无需重复关闭")
                
    def get_status(self) -> Dict[str, Any]:
        """
        获取线程池状态信息
        
        Returns:
            Dict: 包含线程池状态的字典
        """
        with self._lock:
            return {
                "initialized": self.thread_pool is not None,
                "closed": self._closed,
                "max_workers": self.max_workers,
                "thread_name_prefix": self.thread_name_prefix
            }
            
    def __enter__(self):
        """\支持上下文管理器协议"""
        self.initialize()
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        """\支持上下文管理器协议"""
        self.shutdown(wait=True)
        return False
        
# 创建全局线程池管理器实例
# 可以通过配置服务来动态调整参数
thread_pool_manager = ThreadPoolManager(
    max_workers=5,
    thread_name_prefix="rag-"
)